import sys
from numpy import add

from pyparsing import col
sys.path.append('..')
from module import Module
import torch.nn as nn
import pandas as pd


import torch.nn as nn
import json
from collections import defaultdict, OrderedDict
import hashlib

class Module:
    module_map = defaultdict(int)
    src_modules_map = defaultdict(int)
    

    def __init__(self, nn_module: nn.Module, src=False):
        # self.cls = nn_module.__class__
        property = {k: str(v) for k, v in nn_module.__dict__.items() if k not in ['_parameters','_modules', '_buffers'] and v}
        
        self.property = {
            'class': str(nn_module.__class__).split("'")[1],   # class attr
            'type': str(type(nn_module)).split("'")[1],         # ob_type
            **property
        }

        # {k: v for k, v in nn_module.__dict__.items() if not k.startswith('_')}
        self.level = 0
        self.hash = int(hashlib.md5(str(self).encode(encoding='UTF-8')).hexdigest(), 16)
        self.children = {}
        for idx, m in enumerate(nn_module.children()):
            self.children[idx] = Module(m)
            self.hash += self.children[idx].hash << idx
            self.level = max(self.level, self.children[idx].level + 1)

        self.name = self.property['class'].split('.')[-1] + str(self.hash)[-5:]
        Module.module_map[self] += 1
        if src:
            Module.src_modules_map[self] += 1

    def print(self, indent=0):
        print('\t' * indent, self.property)
        for k, v in self.children.items():
            v.print(indent + 1)

    def __hash__(self) -> int:
        return self.hash

    def __eq__(self, othr):
        return (isinstance(othr, type(self))
                and (self.property, self.children) ==
                    (othr.property, othr.children))

    def __str__(self):
        return json.dumps(self.property) 

    def draw(self):
        node_set = set()
        node_set.add(self.name)
        edge_set = set()
        for m in self.children.values():
            ns, es = m.draw()
            node_set |= ns
            edge_set |= es
            edge_set.add((self.name, m.name))
        return node_set, edge_set

    @classmethod
    def statistic(cls, *models):
        from collections import OrderedDict
        import matplotlib.pyplot as plt
        import os
        for x in models:
            cls(x)
        s = []
        s.append(f'module_count: {sum(cls.module_map.values())}')
        s.append(f'unique_count: {len(cls.module_map)}')

        module_dic = Module.sort_map(cls.module_map)
        for k, v in module_dic.items():
            s.append(f'{k.name} (lv{k.level}) cnt:{str(v)} {str(k)}')

        fig, ax = plt.subplots(figsize=(40, 10))
        plt.xticks(rotation=270)
        count = 100
        os.makedirs('output', exist_ok=True)
        
        ax.bar([x.name for x in list(module_dic.keys())[:count]], list(module_dic.values())[:count])
        plt.savefig('output/mydraw.png')

        edge_set = set()
        for m in cls.src_modules_map:
            _, es = m.draw()
            edge_set |= es

        s.append('==========draw_graphgvis==========')
        s.append('digraph G {')
        for e in edge_set:
            s.append(f'  {e[0]} -> {e[1]}')
        for m in cls.src_modules_map:
            s.append(f'  {m.name} [shape=Msquare];')
        s.append('}')
        with open('output/log', 'w') as f:
            f.write('\n'.join(s))

    @staticmethod
    def sort_map(m):
        lis = sorted(m.items(), key=lambda x: x[1], reverse=True)
        dic = OrderedDict()
        for (k, v) in lis:
            dic[k] = v
        return dic

    @classmethod
    def export(cls, *models):
        for x in models:
            cls(x)
        module_dic = cls.sort_map(cls.module_map)

        # 导出modules, prototype
        module_df = pd.DataFrame(columns=['name', 'hash', 'level', 'property', 'prototype'])
        module_df.index.name = "id"
        for x in module_dic:
            # print(x.name, x.property)
            module_df.loc[len(module_df)] = [x.name, x.hash, x.level, x.property, x.property['class'].split('.')[-1]]
        module_df.to_csv('output/module.csv')

        # 导出class, extend
        class_df = pd.DataFrame(columns=['name', 'extend'])
        class_df.index.name = "id"
        class_set = set()
        for x in module_dic:
            class_set.add(x.property['class'])

        cls_set = set()
        def add_class(x):
            if x.count('.'):
                prefix = '.'.join(x.split('.')[:-1])
                x = x.split('.')[-1]
                if x not in cls_set:
                    cls_set.add(x) 
                    class_df.loc[len(class_df)] = [x, add_class(prefix)]
            return x

        for x in class_set:
            add_class(x)
        class_df.to_csv('output/class.csv')

        # 导出construct
        construct_df = pd.DataFrame(columns=['src', 'dest', 'order'])
        construct_df.index.name = "id"
        for x in module_dic:
            for ix, c in x.children.items():
                construct_df.loc[len(construct_df)] = [c.name, x.name, ix]
        construct_df.to_csv('output/construct.csv')


if __name__ == "__main__":
    from torchvision.models import *
    models = [
        resnet18(),
        resnet34(),
        resnet101(),
        resnet152(),
        googlenet(),
        vgg11_bn(),
        vgg13_bn(),
        vgg16(),
        vgg19(),
        squeezenet1_0(),
        squeezenet1_1()
    ]
    Module.statistic(*models)
    Module.export()